import math
import numpy as np
import tensorflow as tf

from tensorflow.python.framework import ops

try:
  image_summary = tf.image_summary
  scalar_summary = tf.scalar_summary
  histogram_summary = tf.histogram_summary
  merge_summary = tf.merge_summary
  SummaryWriter = tf.train.SummaryWriter
except:
  image_summary = tf.summary.image
  scalar_summary = tf.summary.scalar
  histogram_summary = tf.summary.histogram
  merge_summary = tf.summary.merge
  SummaryWriter = tf.summary.FileWriter

if "concat_v2" in dir(tf):
  def concat(tensors, axis, *args, **kwargs):
    return tf.concat_v2(tensors, axis, *args, **kwargs)
else:
  def concat(tensors, axis, *args, **kwargs):
    return tf.concat(tensors, axis, *args, **kwargs)

class batch_norm(object):
  def __init__(self, epsilon=1e-5, momentum = 0.9, name="batch_norm"):
    with tf.variable_scope(name):
      self.epsilon  = epsilon
      self.momentum = momentum
      self.name = name

  def __call__(self, x, train=True):
    return tf.contrib.layers.batch_norm(x,
                      decay=self.momentum,
                      updates_collections=None,
                      epsilon=self.epsilon,
                      scale=True,
                      is_training=train,
                      scope=self.name)

def conv_cond_concat(x, y):
  """Concatenate conditioning vector on feature map axis."""
  x_shapes = x.get_shape()
  y_shapes = y.get_shape()
  return concat([
    x, y*tf.ones([x_shapes[0], x_shapes[1], x_shapes[2], y_shapes[3]])], 3)

def conv3d(input_, out_channels, f_d = 5, f_h = 5, f_w = 5,
           s_d = 2, s_h = 2, s_w = 2, stddev = 0.02, name = 'conv3d'):
    '''
    f - filter
    s - stride
    h - height
    d - depth
    w - width
    '''
    with tf.variable_scope(name):
        filters = tf.get_variable('filters', [f_d, f_h, f_w, input_.get_shape()[-1], out_channels],
                                  initializer=tf.truncated_normal_initializer(stddev=stddev))
        conv = tf.nn.conv3d(input_, filters, strides=[1, s_d, s_h, s_w, 1], padding='SAME')

        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
        conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())

        return conv

def deconv3d(input_, out_shape, f_d = 5, f_h = 5, f_w = 5,
             s_d = 2, s_h = 2, s_w = 2, stddev=0.02, name = 'deconv3d'):
    with tf.variable_scope(name):
        filters = tf.get_variable('filters', [f_d, f_h, f_w, out_shape[-1], input_.get_shape()[-1]],
                                  initializer=tf.random_normal_initializer(stddev=stddev))
        deconv = tf.nn.conv3d_transpose(input_, filters, output_shape=out_shape,
                                        strides=[1, s_d, s_h, s_w, 1])

        biases = tf.get_variable('biases', [out_shape[-1]], initializer=tf.constant_initializer(0.0))
        deconv = tf.reshape(tf.nn.bias_add(deconv, biases), deconv.get_shape())

        return deconv

def dilated_conv3d(inputs, f_d = 3, f_h = 3, f_w = 3,
                 dilation = 2, name = 'dilatedConv3d'):
    '''
    inputs: A tensor with shape [batches, depth, height, width, in_channels]

    return: A tensor has the same shape with 'inputs'
    '''
    with tf.variable_scope(name):
        filters = tf.get_variable('filters', [f_d, f_h, f_w, inputs.get_shape()[-1] , inputs.get_shape()[-1]],
                                  initializer=tf.random_normal_initializer(stddev=0.02))
        # filters with shape [depth, height, width, in_channels, out_channels ]

        return tf.nn.convolution(input = inputs, filter = filters, padding = 'SAME',
                                 dilation_rate = np.broadcast_to(dilation, (3,)), name = name)

def lrelu(x, leak=0.2, name="lrelu"):
  return tf.maximum(x, leak*x)

def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, with_w=False):
  shape = input_.get_shape().as_list()

  with tf.variable_scope(scope or "Linear"):
    matrix = tf.get_variable("Matrix", [shape[1], output_size], tf.float32,
                 tf.random_normal_initializer(stddev=stddev))
    bias = tf.get_variable("bias", [output_size],
      initializer=tf.constant_initializer(bias_start))
    if with_w:
      return tf.matmul(input_, matrix) + bias, matrix, bias
    else:
      return tf.matmul(input_, matrix) + bias
